import torch

def make_train_state(args):
    return {'stop_early':False,
            'early_stopping_step':0,
            'early_stopping_best_val':1e8,
            'learning_rate':args.learning_rate,
            'epoch_index':0,
            'train_loss':[],
            'train_acc':[],
            'val_loss':[],
            'val_acc':[],
            'test_loss':-1,
            'test_acc':-1,
            'model_filename':args.model_state_file}

#返回训练更新梯度的状态
def update_train_state(args,model,train_state):
        # 开始时保存一次模型
        if train_state['epoch_index']==0:
                torch.save(model.state_dict(),train_state['model_filename'])
                train_state['stop_early']=False
        #损失函数减小，保存模型
        elif train_state['epoch_index']>=1:
                loss_tm1,loss_t=train_state['val_loss'][-2:]
                #判断是否大于阈值
                if loss_t>=train_state['early_stopping_best_val']:
                        train_state['early_stopping_step']+=1
                else:
                        #如果小于阈值，保存模型
                        if loss_t<train_state['early_stopping_best_val']:
                                torch.save(model.state_dict(),train_state['model_filename'])
                        #重置
                        train_state['early_stopping_step']=0

                train_state['stop_early']=train_state['early_stopping_step']>=args.early_stopping_criteria

        return train_state

#计算准确率
def compute_accuracy(y_pred,y_target):
        _,y_pred_indices=y_pred.max(dim=1)
        n_correct=torch.eq(y_pred_indices,y_target).sum().item()
        return n_correct/len(y_pred_indices)*100

